import sqlite3
from sqlite3 import Connection
import json
from datetime import datetime


class SMTDatabase:
    def __init__(self, db_path="smt_times.db"):
        self.db_path = db_path
        self.connect()
        self.create_tables()

    def connect(self):
        """连接到SQLite数据库"""
        self.conn = sqlite3.connect(self.db_path)
        self.conn.row_factory = sqlite3.Row  # 使返回的行像字典一样

    def create_tables(self):
        """创建数据库表"""
        with self.conn:
            # 创建表
            self.conn.executescript("""
                CREATE TABLE IF NOT EXISTS smt_lines (
                    line_id INTEGER PRIMARY KEY AUTOINCREMENT,
                    line_name TEXT NOT NULL,
                    description TEXT,
                    is_active INTEGER DEFAULT 1,
                    process_config TEXT
                );
                
                CREATE TABLE IF NOT EXISTS processes (
                    process_id INTEGER PRIMARY KEY AUTOINCREMENT,
                    process_code TEXT NOT NULL UNIQUE,
                    process_name TEXT NOT NULL,
                    description TEXT
                );
                
                CREATE TABLE IF NOT EXISTS products (
                    product_id TEXT PRIMARY KEY,
                    product_name TEXT NOT NULL,
                    description TEXT,
                    default_line_id INTEGER,
                    FOREIGN KEY (default_line_id) REFERENCES smt_lines(line_id)
                );
                
                CREATE TABLE IF NOT EXISTS standard_times (
                    time_id INTEGER PRIMARY KEY AUTOINCREMENT,
                    product_id TEXT NOT NULL,
                    line_id INTEGER NOT NULL,
                    process_id INTEGER NOT NULL,
                    standard_time_seconds REAL NOT NULL,
                    created_date TEXT DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE (product_id, line_id, process_id),
                    FOREIGN KEY (product_id) REFERENCES products(product_id),
                    FOREIGN KEY (line_id) REFERENCES smt_lines(line_id),
                    FOREIGN KEY (process_id) REFERENCES processes(process_id)
                );
            """)

            # 插入标准工序数据（如果不存在）
            processes = [
                ('LOAD', '前置放板'),
                ('SPI', '锡膏印刷'),
                ('SPI_INSP', '印刷检测'),
                ('SMT', '贴片机'),
                ('AOI_2D', '2D光学检测'),
                ('REFLOW', '回焊炉'),
                ('AOI_3D', '3D光学检测')
            ]

            for process_code, process_name in processes:
                self.conn.execute(
                    "INSERT OR IGNORE INTO processes (process_code, process_name) VALUES (?, ?)",
                    (process_code, process_name)
                )

    def create_or_update_line(self, line_name, description=None, is_active=True, process_config=None):
        """创建或更新线体及其工序配置"""
        if process_config is None:
            process_config = {"processes": [1, 2, 3, 4, 5, 6, 7]}

        process_config_str = json.dumps(process_config)

        # 检查线体是否已存在
        cursor = self.conn.execute(
            "SELECT line_id FROM smt_lines WHERE line_name = ?",
            (line_name,)
        )
        existing_line = cursor.fetchone()

        if existing_line:
            # 更新现有线体
            line_id = existing_line['line_id']
            self.conn.execute(
                "UPDATE smt_lines SET description = ?, is_active = ?, process_config = ? WHERE line_id = ?",
                (description, 1 if is_active else 0, process_config_str, line_id)
            )
        else:
            # 创建新线体
            cursor = self.conn.execute(
                "INSERT INTO smt_lines (line_name, description, is_active, process_config) VALUES (?, ?, ?, ?)",
                (line_name, description, 1 if is_active else 0, process_config_str)
            )
            line_id = cursor.lastrowid

        self.conn.commit()
        return line_id

    def create_or_update_product(self, product_id, product_name, description=None, default_line_name=None):
        """创建或更新产品及其默认线体"""
        default_line_id = None

        if default_line_name:
            # 获取或创建默认线体
            default_line_id = self.create_or_update_line(default_line_name)

        # 创建或更新产品
        self.conn.execute(
            """INSERT INTO products (product_id, product_name, description, default_line_id)
               VALUES (?, ?, ?, ?)
               ON CONFLICT(product_id) DO UPDATE SET
               product_name = excluded.product_name,
               description = excluded.description,
               default_line_id = excluded.default_line_id""",
            (product_id, product_name, description, default_line_id)
        )

        self.conn.commit()
        return product_id

    def add_standard_time(self, product_id, line_name, process_code, standard_time_seconds):
        """为产品添加标准工时数据"""
        # 获取线体ID
        cursor = self.conn.execute(
            "SELECT line_id FROM smt_lines WHERE line_name = ?",
            (line_name,)
        )
        line = cursor.fetchone()
        if not line:
            raise ValueError(f"线体 '{line_name}' 不存在")
        line_id = line['line_id']

        # 获取工序ID
        cursor = self.conn.execute(
            "SELECT process_id FROM processes WHERE process_code = ?",
            (process_code,)
        )
        process = cursor.fetchone()
        if not process:
            raise ValueError(f"工序 '{process_code}' 不存在")
        process_id = process['process_id']

        # 插入或更新标准工时
        self.conn.execute(
            """INSERT OR REPLACE INTO standard_times 
               (product_id, line_id, process_id, standard_time_seconds)
               VALUES (?, ?, ?, ?)""",
            (product_id, line_id, process_id, standard_time_seconds)
        )

        self.conn.commit()
        return self.conn.total_changes  # 返回受影响的行数

    def get_all_times(self):
        """获取所有产品的SMT各线工时"""
        cursor = self.conn.execute("""
            SELECT 
                p.product_id,
                p.product_name,
                l.line_name,
                MAX(CASE WHEN pr.process_code = 'LOAD' THEN st.standard_time_seconds END) AS load_time,
                MAX(CASE WHEN pr.process_code = 'SPI' THEN st.standard_time_seconds END) AS spi_time,
                MAX(CASE WHEN pr.process_code = 'SPI_INSP' THEN st.standard_time_seconds END) AS spi_insp_time,
                MAX(CASE WHEN pr.process_code = 'SMT' THEN st.standard_time_seconds END) AS smt_time,
                MAX(CASE WHEN pr.process_code = 'AOI_2D' THEN st.standard_time_seconds END) AS aoi_2d_time,
                MAX(CASE WHEN pr.process_code = 'REFLOW' THEN st.standard_time_seconds END) AS reflow_time,
                MAX(CASE WHEN pr.process_code = 'AOI_3D' THEN st.standard_time_seconds END) AS aoi_3d_time,
                MAX(st.created_date) AS created_date
            FROM products p
            CROSS JOIN smt_lines l
            LEFT JOIN standard_times st ON p.product_id = st.product_id AND l.line_id = st.line_id
            LEFT JOIN processes pr ON st.process_id = pr.process_id
            GROUP BY p.product_id, p.product_name, l.line_name
            ORDER BY p.product_id, l.line_name
        """)

        return cursor.fetchall()

    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


# 使用示例
if __name__ == "__main__":
    # 创建数据库连接
    with SMTDatabase() as db:
        # 创建线体
        db.create_or_update_line(
            "SMT-01",
            "一号SMT生产线，配备高速贴片机",
            True,
            {"processes": [1, 2, 3, 4, 5, 6, 7], "machine_counts": {"1": 1, "2": 1, "3": 1, "4": 3, "5": 1, "6": 1, "7": 1}}
        )

        db.create_or_update_line(
            "SMT-02",
            "二号SMT生产线，配备精密贴片机",
            True,
            {"processes": [1, 2, 4, 6], "machine_counts": {"1": 1, "2": 1, "4": 2, "6": 1}}
        )

        # 创建产品
        db.create_or_update_product(
            "6708-123456-01",
            "主板A型",
            "高性能主板",
            "SMT-01"
        )

        db.create_or_update_product(
            "6708-234567-01",
            "显示模块V1",
            "高清显示模块",
            "SMT-02"
        )

        # 添加标准工时数据
        times_data = [
            # 产品6708-123456-01在SMT-01线的工时
            ("6708-123456-01", "SMT-01", "LOAD", 120.50),
            ("6708-123456-01", "SMT-01", "SPI", 180.25),
            ("6708-123456-01", "SMT-01", "SPI_INSP", 90.75),
            ("6708-123456-01", "SMT-01", "SMT", 360.00),
            ("6708-123456-01", "SMT-01", "AOI_2D", 60.50),
            ("6708-123456-01", "SMT-01", "REFLOW", 240.25),
            ("6708-123456-01", "SMT-01", "AOI_3D", 120.75),

            # 产品6708-123456-01在SMT-02线的工时
            ("6708-123456-01", "SMT-02", "LOAD", 110.50),
            ("6708-123456-01", "SMT-02", "SPI", 170.25),
            ("6708-123456-01", "SMT-02", "SPI_INSP", 85.75),
            ("6708-123456-01", "SMT-02", "SMT", 340.00),
            ("6708-123456-01", "SMT-02", "AOI_2D", 55.50),
            ("6708-123456-01", "SMT-02", "REFLOW", 230.25),
            ("6708-123456-01", "SMT-02", "AOI_3D", 115.75),

            # 产品6708-234567-01在SMT-02线的工时
            ("6708-234567-01", "SMT-02", "LOAD", 100.50),
            ("6708-234567-01", "SMT-02", "SPI", 150.25),
            ("6708-234567-01", "SMT-02", "SPI_INSP", 80.75),
            ("6708-234567-01", "SMT-02", "SMT", 300.00),
            ("6708-234567-01", "SMT-02", "AOI_2D", 50.50),
            ("6708-234567-01", "SMT-02", "REFLOW", 200.25),
            ("6708-234567-01", "SMT-02", "AOI_3D", 100.75),
        ]

        for product_id, line_name, process_code, time_seconds in times_data:
            db.add_standard_time(product_id, line_name, process_code, time_seconds)

        # 查询并显示所有工时数据
        results = db.get_all_times()
        print("产品ID\t\t产品名称\t\t线体\t前置放板\t印刷\t印刷检测\tSMT\t2D检测\t回焊炉\t3D检测\t建立日期")
        print("-" * 120)

        for row in results:
            print(f"{row['product_id']}\t{row['product_name']}\t{row['line_name']}\t"
                  f"{row['load_time'] or 'N/A'}\t{row['spi_time'] or 'N/A'}\t{row['spi_insp_time'] or 'N/A'}\t"
                  f"{row['smt_time'] or 'N/A'}\t{row['aoi_2d_time'] or 'N/A'}\t{row['reflow_time'] or 'N/A'}\t"
                  f"{row['aoi_3d_time'] or 'N/A'}\t{row['created_date'] or 'N/A'}")
